from collections.abc import Generator
import os
import random
import sys
import codecs
import time
import traceback
from sqlmodel import SQLModel
from termcolor import colored
from LLMService import LLMService,PrettyMessages
from classDef import SOURCE_ROOT,reinitDB,affmap,escapeLn,cache_emb,engine
from crudAgent import crudAgent
from crudUser import crudUser
from treeKernel import ftree
import json
from crudAgent import AGENT_TOOLS
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
import warnings
import shlex
warnings.filterwarnings("ignore", category=DeprecationWarning)
from classDef import docker_init
docker_init()
#https://platform.openai.com/docs/overview
sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, errors='ignore')#type: ignore

class NoteAgent:
    def __init__(self):
        self.tree = ftree()
        self.messages = []
        self.LLMservice = LLMService()
        self.crudagent = crudAgent(self.tree)
        self.cruduser = crudUser(self.tree)
        self.curuser= "root"
    def stepPre(self, task: str) -> None:
        # 如下是现有文件目录信息以及根据任务的关键词自动搜索到的相关结果\n{ragstr}.
        self.messages.clear()
        self.messages.append({
            "role":"system",
            "content":f'''你是一个Markdown文件管理助手。
        请根据用户的任务，决定是否新建文件/删除文件/查看某个文件全部内容/继续搜素某个关键词/如果你已经看到原文或文章本身就为空，则可以调用editDocModel修改文章.
        你可以一次性调用多个函数，它们会依次执行。如果你不返回任何函数，那么代表任务终止.
        如果任务的success字段是小于0的值，那么代表调用失败'''
        })
        print(PrettyMessages(self.messages))#system
        self.messages.append({
            "role":"user",
            "content":f"Now call functions. Remember, the task is {task}"
        })
        print(PrettyMessages(self.messages))#user
        sys.stdout.flush()
    


    def stepMain(self,user:str,task:str):
        # 调用 filterMessagesForOpenAI() 过滤掉所有 tool 消息，
        # 再交给 LLMService.funcCallLLMS1()
        reply = self.LLMservice.funcCallLLMS1(
            "",
            AGENT_TOOLS,
            self.messages
        )
        self.messages.append(reply)
        print(PrettyMessages(self.messages))
        sys.stdout.flush()
        # 如果有工具调用
        if reply.tool_calls:
            for tool_call in reply.tool_calls:
                name = tool_call.function.name
                args = json.loads(tool_call.function.arguments)
                tool_result = self.crudagent.router(user, name, args)
                self.addPntToolRtn(tool_result,tool_call)
            return self.stepMain(user,task)
        return

    def permissionBreak(self,user:str,tool_result_store:list[dict],tool_call:ChatCompletionMessageToolCall)->Generator[bytes,None,None]:
        node=self.tree._findNodeById(abs(int(tool_result_store[0].get("success",999999999999))))
        if node:
            self.tree.pendingPermissions[(user,node.id)]="pending"
        else:
            tool_result_store[0]["message"]=f"node {abs(int(tool_result_store[0].get('success',-1)))} not found"
            return
        stream_pr={
            "type":"permission",
            "content":f"Do you agree agent {tool_call.function.name} on \n{self.tree.pNode(node,style='ja')}?",
            "nodeid":node.id
        }
        yield json.dumps(stream_pr).encode("utf-8")
        startTime=time.time()
        while True:
            time.sleep(5)
            print(colored(f"pendingPermissions inside: {self.tree.pendingPermissions}","blue"))
            if time.time()-startTime>100:
                self.tree.pendingPermissions[(user,node.id)]="rejected"
                return
            if self.tree.pendingPermissions.get((user,node.id),"") in ["resolved","rejected"]:
                tool_result_store[0]=self.crudagent.router(user, tool_call.function.name, json.loads(tool_call.function.arguments))
                break
    def addPntToolRtn(self,tool_result:dict, tool_call: ChatCompletionMessageToolCall)->None:
        """
        # 注意，只会push toolcall result ，不会push toolcall到self.messages
        执行工具调用并构造响应消息
        
        Args:
            user: 用户名
            tool_name: 工具名称 
            tool_args: 工具参数
            tool_id: 工具调用ID
        Returns:
            构造的工具响应消息
        """
        tool_result["tool_name"] = tool_call.function.name
        tool_message = {
            "role": "tool",
            "tool_call_id": tool_call.id, 
            "content": json.dumps(tool_result)
        }
        self.messages.append(tool_message)
        print(PrettyMessages(self.messages))#toolreturn
        sys.stdout.flush()
        return 
    def stepMainS(self,user:str,task:str,streamType:int):
        try:
            assistant_content=""
            final_tool_calls = {}
            if streamType==1:
                reply = self.LLMservice.funcCallLLMS1("",AGENT_TOOLS,self.messages)
                assistant_content=reply.content or ""
                for i,tool_call in enumerate(reply.tool_calls or []):
                    final_tool_calls[i] = tool_call
                if reply.tool_calls:
                    stream_tc ={
                        "type": "toolcalls",
                        "toolcalls": [{
                            "id": tool.id,
                            "function": { 
                                "name": tool.function.name,
                                "arguments": tool.function.arguments
                            }
                            } for tool in reply.tool_calls],
                        "content": reply.content,
                    }
                    yield json.dumps(stream_tc).encode("utf-8")
                else:
                    stream_as = {"type": "assistant","content": reply.content or ""}
                    yield json.dumps(stream_as).encode("utf-8")
            elif streamType==2:
                stream = self.LLMservice.funcCallLLMS2( "",AGENT_TOOLS,self.messages)
                replyId=str(random.randint(0,1000000))
                for chunk in stream:
                    # time.sleep(0.5)
                    if not chunk.choices:
                        continue
                    for tool_call in chunk.choices[0].delta.tool_calls or []:
                        index = tool_call.index
                        if index not in final_tool_calls:
                            final_tool_calls[index] = tool_call
                        final_tool_calls[index].function.arguments += tool_call.function.arguments
                        stream_tc_delta={
                            "type":"toolcalls",
                            "toolcalls":[{
                                "id":final_tool_calls[index].id,
                                "function":{
                                    "name":final_tool_calls[index].function.name,
                                    "arguments":tool_call.function.arguments
                                }
                            }],
                            "deltaid":final_tool_calls[index].id
                        }
                        print(stream_tc_delta)
                        sys.stdout.flush()       
                        yield json.dumps(stream_tc_delta).encode("utf-8")
                    if chunk.choices[0].delta.content:
                        assistant_chunk=chunk.choices[0].delta.content or ""
                        assistant_content+=assistant_chunk
                        stream_as_delta={
                            "type":"assistant",
                            "content":assistant_chunk,
                            "deltaid":replyId
                        }
                        print(stream_as_delta)
                        sys.stdout.flush()
                        yield json.dumps(stream_as_delta).encode("utf-8")
            else:
                raise ValueError(f"Invalid streamType: {streamType}")
            self.messages.append({"role":"assistant","tool_calls":[ tool_call for tool_call in final_tool_calls.values() ],"content":assistant_content})
            print(PrettyMessages(self.messages))#assistant/toolcall
            sys.stdout.flush()
            if final_tool_calls:
                for tool_call in final_tool_calls.values():
                    tool_result_store = [self.crudagent.router(user, tool_call.function.name, json.loads(tool_call.function.arguments))]
                    cnt=0
                    print(colored(f"pendingPermissions: {self.tree.pendingPermissions}","cyan"))
                    print(colored(f"tool_result_store: {tool_result_store}","cyan"))
                    sys.stdout.flush()
                    lastPermissionPending=self.tree.pendingPermissions.copy()
                    while "PERMISSION_IN_THE_LOOP" in tool_result_store[0].get("message",""):
                        yield from self.permissionBreak(user,tool_result_store,tool_call)
                        cnt+=1
                        print(colored(f"{cnt} pendingPermissions: {self.tree.pendingPermissions}","cyan"))
                        print(colored(f"{cnt} tool_result_store: {tool_result_store}","cyan"))
                        sys.stdout.flush()
                    self.tree.pendingPermissions=lastPermissionPending
                    self.addPntToolRtn(tool_result_store[0],tool_call)
                    stream_tr={
                        "type":"toolreturn",
                        "toolreturn":{
                            "id":tool_call.id,
                            "function":{
                                "name":tool_call.function.name,
                                "result":tool_result_store[0]
                            }
                        }
                    }
                    print(len(tool_result_store[0].get("nodes",[])))
                    sys.stdout.flush()
                    yield json.dumps(stream_tr).encode("utf-8")
                yield from self.stepMainS(user,task,streamType)
                yield json.dumps({"type":"end"}).encode("utf-8")
                print("STREAMING_END")
                sys.stdout.flush()
            return
        except Exception as e:
            stream_error = {
                "content": traceback.format_exc(),
                "type": "error"
            }
            yield json.dumps(stream_error).encode("utf-8")
            print(f"Error in stepMainS: {traceback.format_exc()}")
            return


    def main(self,user:str,task:str):
        self.stepPre(task)
        self.stepMain(user,task)
    def mainS(self, user: str, task: str,streamType:int):
        self.stepPre(task)
        yield from self.stepMainS(user, task,streamType)
    def login(self):
        while True:
            user_input = shlex.split(input("login as user (or 'exit' to quit): ").strip(), posix=False)
            if not user_input:
                continue
            if user_input[0].lower() == "exit":
                exit()
            self.curuser = user_input[0]
            result=self.tree.addUser(self.curuser)
            print(f"add user result: {escapeLn(result)}")
            break
    def interact(self):
        self.login()
        while True:
            rawinput=input(f"{self.curuser}>>> ")
            query = shlex.split(rawinput.strip(), posix=False)
            if not query:
                continue
            if query[0].lower() == "help":
                for q0 in sorted(["login","ap","share","affmap","import","touch","mkdir","mv","rn","rm","rmdir","cp","cat","rag","ls","agent"]):
                    self.processCLI([q0],rawinput)
            elif query[0].lower() == "sudo":
                self.processCLI(["login","root"],rawinput)
            else:
                self.processCLI(query,rawinput)

    def paramOption(self, enums: list[str], args: list[str],defaultOption:str) -> tuple[str, list[str]]:
        """处理互斥的命令行选项
        Args:
            enums: 可选的参数列表,如["-a","-u"]
            args: 原始参数列表
        Returns:
            (选中的选项,清理后的参数列表)
        """
        selected = defaultOption
        for enum in enums:
            if enum in args:
                selected = enum
                args = [arg for arg in args if arg != enum]
                break
        return selected, args
    def paramPair(self, flag: str, args: list[str],defaultValue:str) -> tuple[str, list[str]]:
        """处理带值的命令行参数对
        Args:
            flag: 要查找的参数标志，如"-sc"
            args: 原始参数列表
        Returns:
            (参数值, 清理后的参数列表)
        """
        if flag in args:
            flag_index = args.index(flag)
            if flag_index + 1 < len(args):
                value = args[flag_index + 1]
                args = [arg for i, arg in enumerate(args) if i != flag_index and i != flag_index + 1]
                return value, args
        return defaultValue, args
    def processCLI(self,query:list[str],rawinput:str):
        cmd = query[0].lower()
        args = query[1:]
        if cmd == "exit":
            exit()
        if cmd == "login":
            if len(args) < 1:
                print("login [newUser]")
                print("login root == sudo")
                return
            self.curuser = args[0]
            result=self.tree.addUser(self.curuser)
            print(f"add user result: {escapeLn(result)}")
            return
        if cmd == "del":
            if len(args)<1:
                print("del [curUser]")
                return
            if args[0] != self.curuser:
                print("only current user can be deleted")
                return
            result=self.tree.delUser(args[0])
            print(f"del user result: {escapeLn(result)}")
            self.login()
            return
        if cmd == "ap":
            if len(args) < 3 or (args[1] not in ["O", "A", "W", "R"]) or not str(args[2]).isdigit():
                print("ap [method: add, override, clear] [operation: O, A, W, R] [node_id]")
                return
            method = args[0]
            operation = args[1]
            result = self.tree.agentPerm(self.curuser, int(args[2]), operation, method)
            self.processCLI(["ls","-u"],"ls -u")
            print(f"set agent permission result: {escapeLn(result)}")
            return
        if cmd == "share": 
            if len(args) < 3:
                print("share [doc/spaceID] [shareUser] [O/A/W/R]")
                return
            result = self.tree.shareNode(self.curuser, args[0], args[2], args[1])
            self.processCLI(["ls","-u"],"ls -u")
            print(f"share: {escapeLn(result)}")
            return
        if cmd == "affmap":
            if len(args) < 1:
                print("affmap [-a] [-d mapID]")
                return
            if args[0] == "-a":
                affmapList = self.tree.session.query(affmap).all()
                for item in affmapList:
                    print(item)
                return
            if args[0]=="-d" and len(args)>1 and str(args[1]).isdigit():
                if input(f"confirm to delete affmap {args[1]}? (y/n): ").lower() == "y":
                    delMap = self.tree.session.query(affmap).where(affmap.intId==int(args[1])).first()#type: ignore
                    print(delMap)
                    self.tree.session.delete(delMap)
                    self.tree.session.commit()
                    print(f"affmap {delMap.intId} deleted")
            return
        if cmd == "ce":
            if len(args) < 1 or (args[0] not in ["-a","-d"]):
                print("ce [-a|-d]")
                return
            if args[0]=="-a":
                cacheEmbList = self.tree.session.query(cache_emb).all()
                for i,item in enumerate(cacheEmbList):
                    print(f"{i}: {item.content[:70]}")
            if args[0]=="-d":
                if input("confirm to delete all cache_emb? (y/n): ").lower() == "y":
                    self.tree.session.query(cache_emb).delete()
                    self.tree.session.commit()
                    print("deleted all cache_emb")
            return
        if cmd == "import":
            if len(args) < 2 or (not str(args[0]).isdigit()):
                print(r"import [spaceID] [path\to\markdown\file]")
                return
            fileName=os.path.basename(args[1])
            print(f"importing {fileName} to space {args[0]}")
            docDict=self.crudagent.createDoc(self.curuser, args[0], fileName)
            if docDict.get("success",-1) < 0:
                print(colored(f"Document import failed: {escapeLn(docDict)}","red"))
                return
            with open(args[1], "r", encoding="utf-8") as file:
                content = file.read()
            result = self.crudagent.writeDoc(self.curuser, docDict.get("success",-1), content)
            print(f"Document import: {escapeLn(result)}")
            return
        # 符合Linux惯例的文档管理命令
        if cmd == "touch":  # Create document
            if len(args) < 2:
                print("touch [spaceID] [fileName]")
                return
            result = self.crudagent.createDoc(self.curuser, args[0], args[1])
            print(f"Document created: {escapeLn(result)}")
            return
        if cmd == "mkdir":  # Create space
            if len(args) < 1:
                print("mkdir [dirName]")
                return
            result = self.crudagent.createSpace(self.curuser, args[0])
            print(f"Space create: {escapeLn(result)}")
            return
        if cmd == "mv":  # Rename/Move
            if len(args) < 2:
                print("mv [srcDocID] [destSpaceID]")
                return
            result = self.crudagent.moveDoc(self.curuser, args[0], args[1])
            print(f"Document moved: {escapeLn(result)}")
            return
        if cmd == "rn":
            if len(args) < 2:
                print("rn [docID] [newName]")
                return
            result = self.crudagent.renameDoc(self.curuser, args[0], args[1])
            print(f"Document rename: {escapeLn(result)}")
            return
        if cmd == "rndir":
            if len(args) < 2:
                print("rndir [dirID] [newName]")
                return
            result = self.crudagent.renameSpace(self.curuser, args[0], args[1])
            print(f"Space rename: {escapeLn(result)}")
            return
        if cmd == "rm":  # Delete document
            if len(args) < 1:
                print("rm [docID]")
                return
            result = self.crudagent.deleteDoc(self.curuser, args[0])
            print(f"Document delete: {escapeLn(result)}")
            return
        if cmd == "rmdir":  # Delete space
            if len(args) < 1:
                print("rmdir [spaceID]")
                return
            result = self.crudagent.deleteSpace(self.curuser, args[0])
            print(f"Space delete: {escapeLn(result)}")
            return
        if cmd == "cp":  # Copy document
            if len(args) < 2:
                print("cp [srcDocID] [destSpaceID]")
                return
            result = self.crudagent.copyDoc(self.curuser, args[0], args[1])
            print(f"Document copy: {escapeLn(result)}")
            return
        if cmd == "cat":
            if len(args) < 1 or (not str(args[0]).isdigit()):
                print("cat [docID] [-l] [-a]")
                return
            # 使用 paramOption 处理 -a 参数
            option, args = self.paramOption(["-a"], args,"")
            by = "agent" if option == "-a" else "user"
            
            # 使用 paramOption 处理 -l 参数
            option, args = self.paramOption(["-l"], args,"")
            lineCnt = option == "-l"
            
            result = self.crudagent.readDoc(self.curuser, args[0], lineCnt=lineCnt, by=by)
            print(f"Document content: {escapeLn(result)}")
            return
        if cmd == "rag":
            if len(args) < 1:
                print("rag [-q query] [-e ensure_full_return_type] [-r root_id] [-m method] [-b by] [-s style] [-sc summaryLen:contentLen]")
                return
            # 参数默认值
            query_str = ""
            ensure_type = ""
            method = "semantic"
            by = "user"
            style = "cat"  # 将 order_by 改为 style
            sub_root_id = 1
            summaryContent="0:50"
            # 解析参数
            i = 0
            while i < len(args):
                if args[i] == '-q':
                    query_str = args[i+1]
                    i += 2
                elif args[i] == '-e':
                    ensure_type = args[i+1]
                    i += 2
                elif args[i] == '-r':
                    sub_root_id = int(args[i+1])
                    i += 2
                elif args[i] == '-m':
                    method = args[i+1]
                    i += 2
                elif args[i] == '-s':  # 将 -o 改为 -s
                    style = args[i+1]  # 将 order_by 改为 style
                    i += 2
                elif args[i] == '-b':
                    by = args[i+1]
                    i += 2
                elif args[i] == '-sc':
                    summaryContent = args[i+1]
                    i += 2
                else:
                    print(f"invalid parameter: {args[i]}")
                    i += 1
            self.tree.RAG(self.curuser, query_str, ensure_type, method, by, style, sub_root_id, summaryContent)
            return
        if cmd == "ls":
            if len(args) < 1:
                print("ls [-a|-u] [id] [-sc summaryLen:contentLen] [-nc]")
                return
            subrootid=1
            summaryContent, args = self.paramPair("-sc", args,"0:50")
            option, args = self.paramOption(["-a","-u"], args,"-u")
            by = "agent" if option == "-a" else "user"
            style, args = self.paramPair("-s", args,"cat")
            if args and str(args[0]).isdigit():
                print(self.tree.pNode(self.tree._findNodeById(int(args[0])),summaryContent,style))
            else:
                self.tree.RAG(self.curuser,"","doc","literal",by,style,subrootid,summaryContent)
            return
        if cmd=="agent":
            if len(args)<1:
                print("agent [task]")
                return
            self.main(self.curuser, rawinput[len("agent "):])
            return
        print(f"invalid command: {cmd}. type 'help' for help.")
        return 


    

if __name__=="__main__":
    agent = NoteAgent()
    agent.interact()
    # stream=agent.mainS("testuser","what is in doc MongoDB? summarize for me",2)
    # for chunk in stream:
    #     # print(chunk)
    #     sys.stdout.flush()
    #     pass
#流程
#dir 的summary只用作RAG emb时，而不发送给LLM，因为就算发送了，dir的summary,其内容是所有子文件summary的简单拼接，却不包含每个文件节点的id，那么对于写文件到哪里的信息没有帮助
#1.用户需求为  写一个内容 -> 先要知道此内容在哪个文件中/有没有关于这个内容的文件存在 -> 如果搜素没搜到，要确保可以看到叶子的dir节点的参数  RAG("")(ensure_full_return_type="doc") -> 写入内容
#2.用户需求为  修改/查询某个知识点   ->RAG(query)  ->修改/查询文件块(包括增删改查)
#3.用户需求为  删除所有含某个知识点的文件   ->RAG(query)返回的是沿途所有块，包括文件和目录，所以也可以实现
#4.用户需求为  重构所有文件分类    ->RAG("")(ensure_full_return_type="h3")   ->  相对定位改动
#5.用户需求为  重写/新写某篇文章    ->RAG("topic")(ensure_full_return_type="doc")   ->在某个文件块下写文章内容 -> LLMmerge
# (一个Assumption:如果要用到ensure_full_return_type参数，那么直接用LLMmerge?）
#LLMmerge用于文章结构大动，而相对定位不方便，应该直接用LLMmerge。->如果p[]生成p*[]这步，是在处理一个file类型的节点，那么就LLMmerge？
#对于p[]生成的p*[]中的每一个p*,
    """
    第一步之前

        对用户任务关键词提取，对每个关键词做RAG。
    第一步
        input：给dir/file级所有块，以及根据用户任务的关键词搜素出的某些相关块
        调用：决定增删查哪篇文章（改也得先查看）
        增删查文章执行
    第二步
        对于增+查+改
        input：查到的整篇内容+第一次搜到的某几个相关块
        调用：写文章+文章id   ->如果还要搜索还可以RAG，返回第二步
        写下文章
    AutoApprove是对操作而言(in setting-ADM)还是文件而言(self memory,true)还是对会话的(overall,one switch)
    context=(token<CONTEXT WINDOW), send one file fulltext.
    context=(token>CONTEXT WINDOW), send multi file for RAG.
    context=NULL, send all for RAG.
    
    CRDTS->原文直接同步
    summary/emb/tree structure
        10 min->indexing(autosave)
        if modification is large,index in advance->Agent生成文章时给出，文章有无被大幅修改的标记
        (before RAG->indexing)
        sync button

    sync由AGENT/HUMAN invoke.
    sync后倒计时取消
    倒计时只能由原文小改动发起
    在RAG之前，要sync

    大小改动区分
    人类改都算小改动？->(除了删除)
    Agent给标记

    q1.indexing->tree compare relative location*1 to find the target node/complete reindexing the whole file*1.2  ->WHETHER AND HOW TO DEFFERENCIATE THESE TWO?
    #q2.h2->h3,non-semantic but syntactically modification.->10 min must correct tree structure, alone with node type.But within 10 min, the tree structure is not stable.?HARD!
    #q3.ancestor repeatedly renewing. If the renew time is discrete, then problem resolved.
    q4=q2 *1.2 reuse summary/emb
        type-tree modification
        content modification
    日程  
    
    1.relative add node OR re-indexing whole file
    2.h2->h3 is not changing the semantics
    """